#pragma once

#include "cuda_fp16.h"
#include "cuda_runtime.h"
#include "macro.h"
#include "tensor_cuda.h"
#include "traits.h"
#include <atomic>
#include <cstdint>
#include <cstring>
#include <functional>
#include <numeric>
#include <vector>

namespace nvinfer1 {

class Tensor {
public:
  struct CudaStreamGuard {
    CudaStreamGuard(cudaStream_t stream) { Tensor::CudaStream() = stream; }
    ~CudaStreamGuard() { Tensor::CudaStream() = 0; }
  };

  using TensorSize = std::vector<int32_t>;

  // constructor
  Tensor(const TensorSize &size, const void *data,
         const DataType dtype = DataType::kFLOAT)
      : size_(size), data_(const_cast<void *>(data)), dtype_(dtype) {}

  Tensor(const void *data, const PluginTensorDesc &desc)
      : data_(const_cast<void *>(data)), dtype_(desc.type) {
    for (auto i = 0; i < desc.dims.nbDims; ++i) {
      size_.emplace_back(desc.dims.d[i]);
    }
  }

  // methods
  const DataType &dtype() const { return dtype_; }

  const size_t &volume() const {
    if (volume_ == static_cast<size_t>(-1)) {
      evaluateVolume();
    }
    return volume_;
  }

  const TensorSize &size() const { return size_; }
  int size(const int index) const {
    if (index < 0) {
      return size_[size_.size() + index];
    } else {
      return size_[index];
    }
  }

  size_t dim() const { return size_.size(); }

  size_t dbytes() const {
    size_t bytes;
    DISPATCH_ALL_TYPES(dtype_, [&]() { bytes = sizeof(scalar_t);});
    return bytes;
  }

  size_t bytes() const { return volume() * dbytes(); }

  Tensor &zeros() {

    if (data_) {
      DISPATCH_ALL_TYPES(dtype_, [&]() {
        cudaMemsetAsync(reinterpret_cast<scalar_t *>(data_), 0, bytes(),
                        CudaStream());
      });
    }

    return *this;
  }

  template <typename T> Tensor &fill(const T fill_value) {
    if (data_) {
      DISPATCH_ALL_TYPES(dtype_, [&]() {
        fillCudaMemAsync<scalar_t>(data_, static_cast<scalar_t>(fill_value),
                                   volume(), CudaStream());
      });
    }

    return *this;
  }

  template <typename T = void> T *data() { return static_cast<T *>(data_); }
  template <typename T = void> const T *data() const {
    return static_cast<const T *>(data_);
  }

  template <typename T = void> const T *tail() const {
    return static_cast<T *>((int8_t *)data_ + bytes());
  }

  static cudaStream_t &CudaStream() {
    static cudaStream_t stream = 0;
    return stream;
  }

protected:
  size_t evaluateVolume() const {
    volume_ = std::accumulate(
        size_.begin(), size_.end(), 1,
        [](const size_t &a, const size_t &b) { return a * b; });
    return volume_;
  }

private:
  TensorSize size_;
  DataType dtype_ = DataType::kFLOAT;
  mutable size_t volume_ = static_cast<size_t>(-1);
  void *data_ = nullptr;
};
} // namespace nvinfer1
